import math

import numpy as np

import paddle
import paddle.fluid.layers as L
from paddle.fluid.dygraph import to_variable

from helm.dynamic.engine.callback import Callback


def calculate_gain(p):
    if p == 0 or p == 1:
        return 0
    return -p * math.log(p) - (1 - p) * math.log(1 - p)


class MixLoss:

    def __init__(self, criterion):
        self.criterion = criterion
        self.lam = None

    def __call__(self, input, target):
        if self.lam:
            y_a, y_b = target
            loss = self.lam * self.criterion(input, y_a) + (1 - self.lam) * self.criterion(input, y_b)
            return loss - calculate_gain(self.lam)
        else:
            return self.criterion(input, target)


def rand_bbox(size, lam):
    w = size[2]
    h = size[3]
    cut_rat = np.sqrt(1. - lam)
    cut_w = np.int(w * cut_rat)
    cut_h = np.int(h * cut_rat)

    cx = np.random.randint(w)
    cy = np.random.randint(h)

    bbx1 = np.clip(cx - cut_w // 2, 0, w)
    bby1 = np.clip(cy - cut_h // 2, 0, h)
    bbx2 = np.clip(cx + cut_w // 2, 0, w)
    bby2 = np.clip(cy + cut_h // 2, 0, h)

    return bbx1, bby1, bbx2, bby2


class CutMix(Callback):

    def __init__(self, beta, prob):
        super().__init__()
        self.beta = beta
        self.prob = prob
        self.old_criterion = None

    def on_epoch_begin(self, engine):
        self.old_criterion = engine.criterion
        engine.criterion = MixLoss(engine.criterion)

    def on_batch_begin(self, engine):
        r = np.random.rand()
        if r > self.prob:
            engine.criterion.lam = None
            return

        x, y = engine.state.batch
        lam = np.random.beta(self.beta, self.beta) if self.beta > 0 else 1

        b, c, h, w = x.shape
        index = L.randperm(b)

        y_a, y_b = y, L.index_select(y, index)
        bbx1, bby1, bbx2, bby2 = rand_bbox(x.shape, lam)
        xt = x.numpy()
        xt[:, :, bbx1:bbx2, bby1:bby2] = xt[index.numpy(), :, bbx1:bbx2, bby1:bby2]
        x = to_variable(xt)
        lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (h * w))
        engine.criterion.lam = float(lam)
        engine.state.batch = x, (y_a, y_b)

    def on_epoch_end(self, engine):
        engine.criterion = self.old_criterion


class Mixup(Callback):

    def __init__(self, alpha=1.0):
        super().__init__()
        self.alpha = alpha
        self.old_criterion = None

    def on_epoch_begin(self, engine):
        self.old_criterion = engine.criterion
        engine.criterion = MixLoss(engine.criterion)

    def on_batch_begin(self, engine):

        x, y = engine.state.batch
        lam = np.random.beta(self.alpha, self.alpha) if self.alpha > 0 else 1

        batch_size = x.size()[0]
        index = L.randperm(batch_size)

        x = lam * x + (1 - lam) * L.index_select(x, index)
        y_a, y_b = y, L.index_select(y, index)
        self.lam = lam
        engine.criterion.lam = float(lam)
        engine.state.batch = x, (y_a, y_b)

    def on_epoch_end(self, engine):
        engine.criterion = self.old_criterion


def get_mix(cfg):
    if cfg is None:
        return None
    mix_type = cfg.type
    cfg.pop("type")
    mix = eval(mix_type)(**cfg)
    return mix
